/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File model_process.cpp
* Description: handle model process
*/
#include "dimp/model_process.h"
#include <iostream>
#include "utils.h"
using namespace std;

ModelDimpProcess::ModelDimpProcess():loadFlag_(false), modelId_(0), modelMemPtr_(nullptr), modelMemSize_(0),
modelWeightPtr_(nullptr),modelWeightSize_(0), modelDesc_(nullptr), input_(nullptr), output_(nullptr),
isReleased_(false){

}

ModelDimpProcess::~ModelDimpProcess() {
    DestroyResource();
}

void ModelDimpProcess::DestroyResource() {
    if (isReleased_)
        return;
    
    Unload();
    DestroyDesc();
    DestroyInput();
    DestroyOutput();
    isReleased_ = true;
}

Result ModelDimpProcess::LoadModelFromFileWithMem(const char *modelPath) {
    if (loadFlag_) {
        ERROR_LOG("has already loaded a model");
        return FAILED;
    }

    aclError ret = aclmdlQuerySize(modelPath, &modelMemSize_, &modelWeightSize_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("query model failed, model file is %s", modelPath);
        return FAILED;
    }

    ret = aclrtMalloc(&modelMemPtr_, modelMemSize_, ACL_MEM_MALLOC_HUGE_FIRST);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("malloc buffer for mem failed, require size is %zu", modelMemSize_);
        return FAILED;
    }

    INFO_LOG("start Malloc for weight!");

    ret = aclrtMalloc(&modelWeightPtr_, modelWeightSize_, ACL_MEM_MALLOC_HUGE_FIRST);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("malloc buffer for weight failed, require size is %zu", modelWeightSize_);
        return FAILED;
    }

    INFO_LOG("start LoadFromFileWithMem!");

    ret = aclmdlLoadFromFileWithMem(modelPath, &modelId_, modelMemPtr_,
        modelMemSize_, modelWeightPtr_, modelWeightSize_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("load model from file failed, model file is %s", modelPath);
        return FAILED;
    }

    loadFlag_ = true;
    INFO_LOG("load model %s success", modelPath);
    return SUCCESS;
}

Result ModelDimpProcess::CreateDesc()
{
    modelDesc_ = aclmdlCreateDesc();
    if (modelDesc_ == nullptr) {
        ERROR_LOG("create model description failed");
        return FAILED;
    }

    aclError ret = aclmdlGetDesc(modelDesc_, modelId_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("get model description failed");
        return FAILED;
    }

    INFO_LOG("create model description success");
    return SUCCESS;
}

void ModelDimpProcess::DestroyDesc()
{
    if (modelDesc_ != nullptr) {
        (void)aclmdlDestroyDesc(modelDesc_);
        modelDesc_ = nullptr;
    }
}

Result ModelDimpProcess::Create3Inputs(void *input1, size_t input1size, 
                                 void* input2, size_t input2size,
                                 void* input3, size_t input3size)
{
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        ERROR_LOG("can't create dataset, create input failed");
        return FAILED;
    }

    aclDataBuffer* inputData = aclCreateDataBuffer(input1, input1size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input1 failed");
        return FAILED;
    }


    aclError ret = aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input1 failed");
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData2 = aclCreateDataBuffer(input2, input2size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input2 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData2);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input2 failed");
        aclDestroyDataBuffer(inputData2);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData3 = aclCreateDataBuffer(input3, input3size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input3 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData3);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input3 failed");
        aclDestroyDataBuffer(inputData3);
        inputData = nullptr;
        return FAILED;
    }

    return SUCCESS;
}

Result ModelDimpProcess::Create2Inputs(void *input1, size_t input1size, 
                                 void* input2, size_t input2size)
{
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        ERROR_LOG("can't create dataset, create input failed");
        return FAILED;
    }

    aclDataBuffer* inputData = aclCreateDataBuffer(input1, input1size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input1 failed");
        return FAILED;
    }

    // cout << input1size << endl;

    aclError ret = aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input1 failed");
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData2 = aclCreateDataBuffer(input2, input2size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input2 failed");
        return FAILED;
    }

    // cout << input2size << endl;

    ret = aclmdlAddDatasetBuffer(input_, inputData2);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input2 failed");
        aclDestroyDataBuffer(inputData2);
        inputData = nullptr;
        return FAILED;
    }
    
    return SUCCESS;
}

Result ModelDimpProcess::Create1Inputs(void *input1, size_t input1size)
{
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        ERROR_LOG("can't create dataset, create input failed");
        return FAILED;
    }

    aclDataBuffer* inputData = aclCreateDataBuffer(input1, input1size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input1 failed");
        return FAILED;
    }

    // cout << input1size << endl;

    aclError ret = aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input1 failed");
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return FAILED;
    }
    
    return SUCCESS;
}

Result ModelDimpProcess::Create4Inputs(void *input1, size_t input1size, 
                                 void* input2, size_t input2size, void* input3, size_t input3size,
                                 void* input4, size_t input4size)
{
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        ERROR_LOG("can't create dataset, create input failed");
        return FAILED;
    }

    aclDataBuffer* inputData = aclCreateDataBuffer(input1, input1size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input1 failed");
        return FAILED;
    }

    aclError ret = aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input1 failed");
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData2 = aclCreateDataBuffer(input2, input2size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input2 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData2);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input2 failed");
        aclDestroyDataBuffer(inputData2);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData3 = aclCreateDataBuffer(input3, input3size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input3 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData3);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input3 failed");
        aclDestroyDataBuffer(inputData3);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData4 = aclCreateDataBuffer(input4, input4size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input4 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData4);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input4 failed");
        aclDestroyDataBuffer(inputData4);
        inputData = nullptr;
        return FAILED;
    }
     
    return SUCCESS;
}

Result ModelDimpProcess::Create5Inputs(void *input1, size_t input1size, 
                                 void* input2, size_t input2size, void* input3, size_t input3size,
                                 void* input4, size_t input4size, void* input5, size_t input5size)
{
    input_ = aclmdlCreateDataset();
    if (input_ == nullptr) {
        ERROR_LOG("can't create dataset, create input failed");
        return FAILED;
    }

    aclDataBuffer* inputData = aclCreateDataBuffer(input1, input1size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input1 failed");
        return FAILED;
    }

    aclError ret = aclmdlAddDatasetBuffer(input_, inputData);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input1 failed");
        aclDestroyDataBuffer(inputData);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData2 = aclCreateDataBuffer(input2, input2size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input2 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData2);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input2 failed");
        aclDestroyDataBuffer(inputData2);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData3 = aclCreateDataBuffer(input3, input3size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input3 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData3);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input3 failed");
        aclDestroyDataBuffer(inputData3);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData4 = aclCreateDataBuffer(input4, input4size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input4 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData4);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input4 failed");
        aclDestroyDataBuffer(inputData4);
        inputData = nullptr;
        return FAILED;
    }

    aclDataBuffer* inputData5 = aclCreateDataBuffer(input5, input5size);
    if (inputData == nullptr) {
        ERROR_LOG("can't create data buffer, create input5 failed");
        return FAILED;
    }

    ret = aclmdlAddDatasetBuffer(input_, inputData5);
    if (inputData == nullptr) {
        ERROR_LOG("can't add data buffer, create input5 failed");
        aclDestroyDataBuffer(inputData5);
        inputData = nullptr;
        return FAILED;
    }
     
    return SUCCESS;
}


void ModelDimpProcess::DestroyInput()
{
    if (input_ == nullptr) {
        return;
    }

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(input_); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(input_, i);
        aclDestroyDataBuffer(dataBuffer);
    }
    aclmdlDestroyDataset(input_);
    input_ = nullptr;
}

Result ModelDimpProcess::CreateOutput()
{
    if (modelDesc_ == nullptr) {
        ERROR_LOG("no model description, create ouput failed");
        return FAILED;
    }

    output_ = aclmdlCreateDataset();
    if (output_ == nullptr) {
        ERROR_LOG("can't create dataset, create output failed");
        return FAILED;
    }

    size_t outputSize = aclmdlGetNumOutputs(modelDesc_);
    for (size_t i = 0; i < outputSize; ++i) {
        size_t buffer_size = aclmdlGetOutputSizeByIndex(modelDesc_, i);

        void *outputBuffer = nullptr;
        aclError ret = aclrtMalloc(&outputBuffer, buffer_size, ACL_MEM_MALLOC_NORMAL_ONLY);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("can't malloc buffer, size is %zu, create output failed", buffer_size);
            return FAILED;
        }

        aclDataBuffer* outputData = aclCreateDataBuffer(outputBuffer, buffer_size);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("can't create data buffer, create output failed");
            aclrtFree(outputBuffer);
            return FAILED;
        }

        ret = aclmdlAddDatasetBuffer(output_, outputData);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("can't add data buffer, create output failed");
            aclrtFree(outputBuffer);
            aclDestroyDataBuffer(outputData);
            return FAILED;
        }
    }

    INFO_LOG("create model output success");
    return SUCCESS;
}

void ModelDimpProcess::DestroyOutput()
{
    if (output_ == nullptr) {
        return;
    }

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(output_); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(output_, i);
        void* data = aclGetDataBufferAddr(dataBuffer);
        (void)aclrtFree(data);
        (void)aclDestroyDataBuffer(dataBuffer);
    }

    (void)aclmdlDestroyDataset(output_);
    output_ = nullptr;
}

Result ModelDimpProcess::Execute(aclrtStream now_stream, int execute_way)
{
    // 同步推理
    if (execute_way == 0) {
        
        aclError ret = aclmdlExecute(modelId_, input_, output_);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("execute model failed, modelId is %u", modelId_);
            return FAILED;
        }
    }
    
    if (execute_way == 1) {
        cout<<"yes"<<endl;
        aclError ret = aclmdlExecuteAsync(modelId_, input_, output_, now_stream);       // 异步推理
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("execute model failed, modelId is %u", modelId_);
            return FAILED;
        }
    }
    return SUCCESS;
}



void ModelDimpProcess::Unload()
{
    if (!loadFlag_) {
        WARN_LOG("no model had been loaded, unload failed");
        return;
    }

    aclError ret = aclmdlUnload(modelId_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("unload model failed, modelId is %u", modelId_);
    }

    if (modelDesc_ != nullptr) {
        (void)aclmdlDestroyDesc(modelDesc_);
        modelDesc_ = nullptr;
    }

    if (modelMemPtr_ != nullptr) {
        aclrtFree(modelMemPtr_);
        modelMemPtr_ = nullptr;
        modelMemSize_ = 0;
    }

    if (modelWeightPtr_ != nullptr) {
        aclrtFree(modelWeightPtr_);
        modelWeightPtr_ = nullptr;
        modelWeightSize_ = 0;
    }

    loadFlag_ = false;
    INFO_LOG("unload model success, modelId is %u", modelId_);
}

aclmdlDataset *ModelDimpProcess::GetModelOutputData()
{
    return output_;
}

Result ModelDimpProcess::GetCurOutputDimsMul(size_t index, std::vector<int64_t>& curOutputDimsMul)
{
    aclError ret; 
    aclmdlIODims ioDims;
    int64_t tmp_dim = 1;
    ret = aclmdlGetCurOutputDims(modelDesc_, index, &ioDims);
    if (ret != ACL_SUCCESS) {
        WARN_LOG("aclmdlGetCurOutputDims failed ret[%d], maybe the modle has dynamic shape", ret);
        return FAILED;
    }
    for (int i = 1; i < ioDims.dimCount; i++) {
        tmp_dim *= ioDims.dims[ioDims.dimCount - i];
        curOutputDimsMul.push_back(tmp_dim);
    }
    return SUCCESS;
}

void ModelDimpProcess::OutputModelResult(std::string& modelName)
{
    std::string T = Utils::printCurrentTime();
    const char* temp_s = T.c_str();
    if (NULL == opendir(temp_s)) {
        mkdir(temp_s, 0775);
    }

    void* data = nullptr;
    void* outHostData = nullptr;
    void* outData = nullptr;
    aclError ret = ACL_SUCCESS;    
    uint64_t maxBatchSize = 0;
    size_t len = 0;

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(output_); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(output_, i);
        data = aclGetDataBufferAddr(dataBuffer);
        len = aclGetDataBufferSizeV2(dataBuffer);
        aclDataType datatype = aclmdlGetOutputDataType(modelDesc_, i);
      
        outData = reinterpret_cast<float*>(data);

        vector<int64_t> curOutputDimsMul;
        ret = GetCurOutputDimsMul(i, curOutputDimsMul);
        ofstream outstr(T + "/" + modelName + "_output_" + to_string(i) + ".txt", ios::out);
        if (!outstr.is_open()) {
            cout << "not open"<<endl;
        }
        switch (datatype) {
        case 0:
            for (int64_t i = 1; i <= len / sizeof(float); i++) {
                float out = *((float*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                }                    
            }
            break;
        case 1:{
            aclFloat16 * out_fp16 = reinterpret_cast<aclFloat16*>(outData);
            float out = 0;
            for (int i = 1; i <= len / sizeof(aclFloat16); i++) {
                out = aclFloat16ToFloat(out_fp16[i-1]);

                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                }                   
            }
            break;
        }
        case 2:
            for (int i = 1; i <= len / sizeof(int8_t); i++) {
                int8_t out = *((int8_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 3:
            for (int i = 1; i <= len / sizeof(int); i++) {
                int out = *((int*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                }
            }
            break;
        case 4:
            for (int i = 1; i <= len / sizeof(uint8_t); i++) {
                uint8_t out = *((uint8_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                }
            }
            break;
        case 6:
            for (int i = 1; i <= len / sizeof(int16_t); i++) {
                int16_t out = *((int16_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 7:
            for (int i = 1; i <= len / sizeof(uint16_t); i++) {
                uint16_t out = *((uint16_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 8:
            for (int i = 1; i <= len / sizeof(uint32_t); i++) {
                uint32_t out = *((uint32_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 9:
            for (int i = 1; i <= len / sizeof(int64_t); i++) {
                int64_t out = *((int64_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 10:
            for (int i = 1; i <= len / sizeof(uint64_t); i++) {
                uint64_t out = *((uint64_t*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 11:
            for (int i = 1; i <= len / sizeof(double); i++) {
                double out = *((double*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        case 12:
            for (int i = 1; i <= len / sizeof(bool); i++) {
                int out = *((bool*)outData + i - 1);
                outstr << out << " ";
                vector<int64_t>::iterator it;
                for(it = curOutputDimsMul.begin(); it != curOutputDimsMul.end(); it++){
                    if ((i != 0) && (i % *it == 0)){
                        outstr << "\n";
                        break;
                    } 
                } 
            }
            break;
        default:
            printf("undefined data type!\n");
            break;
        }
        outstr.close();

    }

    INFO_LOG("output data success");
    return;
}


